package edu.zju.gis.dbfg.model.statistic;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.base.Constant;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.*;
import org.apache.spark.broadcast.Broadcast;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Geometry;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.*;

public class PacStatistic {

    private final static Logger logger = LoggerFactory.getLogger(PacStatistic.class);

    public static void main(String[] args) throws Exception {
        // Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
        conf.setMaster("local[4]");
        conf.setAppName("lalala");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        // Double sumArea = 100.0;

        // Setup model arguments
        logger.info("Setup arguments");
        PacStatisticArgs mArgs = PacStatisticArgs.initArgs(args);
        if (mArgs == null) {
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(PacStatistic.class, "main()", msg, args);
        }

        CoordinateReferenceSystem crs1 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs1()));
        CoordinateReferenceSystem targetCRS = Constant.DEFAULT_CRS();
        MathTransform transform1 = null;
        if (!crs1.equals(targetCRS)) {
            transform1 = CRS.findMathTransform(crs1, targetCRS);
        }
        Broadcast<PacStatisticArgs> mArgsBroad = jsc.broadcast(mArgs);

        String input1 = mArgs.getInput();
        String type1 = mArgs.getType();
        String cc1 = mArgs.getCcIndex();
        String pac1 = mArgs.getPacIndex();
        String area1 = mArgs.getAreaIndex();
        Map<String, Object> attrs1 = new HashMap<>();
        attrs1.put("CC", cc1);
        attrs1.put("PAC", pac1);
        if (!area1.equals("-1")) {
            attrs1.put("AREA", area1);
        }
        //attrs1.put("YEAR", mArgs.getYearIndex());
        JavaRDD<GeoItem> gi1 = null;
        List<String> filters = Arrays.asList(mArgs.getFilters().split(","));
        if (type1.equals("ES")) {
            gi1 = DataReaderRDD.readToRDD2(type1, jsc, input1, attrs1, filters, mArgs.getFilterIndex(), transform1);
        } else {
            gi1 = DataReaderRDD.readToRDD(type1, jsc, input1, null, attrs1, transform1);
        }

        JavaPairRDD<String, GeoItem> resultRDD = gi1.mapToPair(new PairFunction<GeoItem, String, GeoItem>() {
            @Override
            public Tuple2<String, GeoItem> call(GeoItem gi1) throws Exception {
                //Geometry g1 = gi1.getGeometry();
                String cc = (String) gi1.getProp().get("CC");
                String pac = String.valueOf(gi1.getProp().get("PAC"));
                String id = pac + "##" + cc;
                //props.put("PAREA", area/sumArea); //TODO 改成实际总面积
                return new Tuple2<>(id, gi1);
            }
        });

        JavaPairRDD<String, GeoItem> rr = resultRDD.reduceByKey(new Function2<GeoItem, GeoItem, GeoItem>() {
                                                                    @Override
                                                                    public GeoItem call(GeoItem gi1, GeoItem gi2) throws Exception {
                                                                        Double area = (Double) gi1.getProp().get("AREA") + (Double) gi2.getProp().get("AREA");
                                                                        //Double pArea = area / sumArea; // TODO 改成实际总面积
                                                                        GeoItem result = new GeoItem();
                                                                        result.setFid(gi1.getFid());
                                                                        result.setProp(gi1.getProp());
                                                                        result.getProp().put("AREA", area);
                                                                        return result;
                                                                    }
                                                                }
        );

        Map<String, GeoItem> result = rr.collectAsMap();
        File file = new File(mArgs.getOutput());
        if (file.exists()) {
            file.delete();
        }
        FileWriter writer = new FileWriter(file);
        if(area1.equals("-1")){
            area1 = "area";
        }
        writer.write(String.format("%s\t%s\t%s\n",pac1, cc1, area1));
        for (String key : result.keySet()) {
            String[] fields = key.split("##");
            String pac = fields[0];
            String cc = fields[1];
            //String scale = fields[1];
            GeoItem gi = result.get(key);
            Double area = (Double) gi.getProp().get("AREA");
            //Double parea = (Double) gi.getProp().get("PAREA");
            String s = String.format("%s\t%s\t%.9f", pac, cc, area);
            writer.write(s + "\n");
        }
        writer.flush();
        writer.close();

    }

    public static void writeToFile(JavaPairRDD<String, GeoItem> rdd, String filename) throws IOException {
        List<Tuple2<String, GeoItem>> featureResult = rdd.collect();
        File featureFile = new File(filename);
        if (featureFile.exists()) featureFile.delete();
        FileWriter featureWriter = new FileWriter(featureFile);
        for (Tuple2<String, GeoItem> feature : featureResult) {
            GeoItem gi = feature._2;
            String id = gi.getFid();
            Geometry geometry = gi.getGeometry();
            String wkt = geometry.toText();
            featureWriter.write(String.format("%s\t%s\n", id, wkt));
        }
        featureWriter.flush();
        featureWriter.close();
    }
}
